clear all
close all
clc
%% Set parameters
% SBGO parameters
num_params = 2;
num_initial_samples = 12;
num_iterations = 4;
% Plotting parameters
plotting.iteration = 4; % SBGO iteration whose pareto points will be plotted: 'all' | 1-num_iterations
plotting.model = 'surrogate'; % Model whose pareto points will be plotted : truth | surrogate
plotting.afs = 18; % Axis Fontsize
plotting.fn = 'Times New Roman'; % Font name
%% Store Pareto Front
% Initial samples
dakota_output = importdata('dakota_surrogate_based_global_optimization.dat');
pareto_front.initial = dakota_output.data(1:num_initial_samples,:); 
% Iteration samples
sample_counter = num_initial_samples;
for i = 1:num_iterations
    filestring = "finaldata" + num2str(i) + ".dat";
    iter_data = importdata(filestring);
    num_iter_samples = size(iter_data,1);
    pareto_front.iter(i).surrogate = importdata(filestring);
    pareto_front.iter(i).truth = dakota_output.data(sample_counter:sample_counter+num_iter_samples,:);
    sample_counter = sample_counter + num_iter_samples;
end
%% Plot 2D
disp('Plotting 2D pareto frontier');
num_iterations = size(pareto_front.iter,2);
pareto_plot = figure(1);
hold on;
% Scaling parameters
a = 10; % 10
b = 25; % 25
% Initial points
a0 = 1;
b0 = size(pareto_front.initial,1);
[~,idx] = sort(pareto_front.initial(:,num_params+3));
for n = 1:size(pareto_front.initial,1)
    scale = (b - a)*(n - a0)/(b0 - a0) + a;
    plot(pareto_front.initial(idx(n),num_params+1),pareto_front.initial(idx(n),num_params+2),'color','k','marker','o','MarkerSize',scale,'LineWidth',2);
end
% Pareto points
if (plotting.iteration == 'all')
    for i = 1:num_iterations
        a0 = 1;
        b0 = size(pareto_front.iter(i).truth,1);
        switch plotting.model
            case 'surrogate'
                [~,idx] = sort(pareto_front.iter(i).surrogate(:,num_params+3));
                for n = 1:size(pareto_front.iter(i).surrogate,1)
                    scale = (b - a)*(n - a0)/(b0 - a0) + a;
                    plot(pareto_front.iter(i).surrogate(idx(n),num_params+1),pareto_front.iter(i).surrogate(idx(n),num_params+2),'color','r','marker','o','MarkerSize',scale,'LineWidth',2);
                end
            case 'truth'
                [~,idx] = sort(pareto_front.iter(i).truth(:,num_params+3));
                for n = 1:size(pareto_front.iter(i).truth,1)
                    scale = (b - a)*(n - a0)/(b0 - a0) + a;
                    plot(pareto_front.iter(i).truth(idx(n),num_params+1),pareto_front.iter(i).truth(idx(n),num_params+2),'color','r','marker','o','MarkerSize',scale,'LineWidth',2);
                end
        end
    end
else
    i = plotting.iteration;
    a0 = 1;
    b0 = size(pareto_front.iter(i).truth,1);
    switch plotting.model
        case 'surrogate'
            [~,idx] = sort(pareto_front.iter(i).surrogate(:,num_params+3));
            for n = 1:size(pareto_front.iter(i).surrogate,1)
                scale = (b - a)*(n - a0)/(b0 - a0) + a;
                plot(pareto_front.iter(i).surrogate(idx(n),num_params+1),pareto_front.iter(i).surrogate(idx(n),num_params+2),'color','r','marker','o','MarkerSize',scale,'LineWidth',2);
            end
        case 'truth'
            [~,idx] = sort(pareto_front.iter(i).truth(:,num_params+3));
            for n = 1:size(pareto_front.iter(i).truth,1)
                scale = (b - a)*(n - a0)/(b0 - a0) + a;
                plot(pareto_front.iter(i).truth(idx(n),num_params+1),pareto_front.iter(i).truth(idx(n),num_params+2),'color','r','marker','o','MarkerSize',scale,'LineWidth',2);
            end
    end
end
hold off;
% Format axes
grid on;
set(pareto_plot,'Color','w');
set(gca,'Fontsize',plotting.afs,'Fontname',plotting.fn);
xlab = xlabel('Averaged Stress - Axial');
set(xlab,'Fontsize',plotting.afs,'Fontname',plotting.fn);
ylab = ylabel('Averaged Stress - Moment');
set(ylab,'Fontsize',plotting.afs,'Fontname',plotting.fn);
% Save as PDF
print(pareto_plot,'pareto2d','-dpdf');
close all
disp('done');
%% Plot 3D
disp('Plotting 3D pareto frontier');
num_iterations = size(pareto_front.iter,2);
pareto_plot = figure(1);
% Initial points
plot3(pareto_front.initial(:,num_params+1),pareto_front.initial(:,num_params+2),pareto_front.initial(:,num_params+3),'ko','LineWidth',2,'MarkerSize',10);
hold on;
% Pareto points
if (plotting.iteration == 'all')
    for i = 1:num_iterations
        switch plotting.model
            case 'surrogate'
                plot3(pareto_front.iter(i).surrogate(:,num_params+1),pareto_front.iter(i).surrogate(:,num_params+2),pareto_front.iter(i).surrogate(:,num_params+3),'color','r','marker','o','LineStyle','none','LineWidth',2,'MarkerSize',10);
                
            case 'truth'
                plot3(pareto_front.iter(i).truth(:,num_params+1),pareto_front.iter(i).truth(:,num_params+2),pareto_front.iter(i).truth(:,num_params+3),'color','r','marker','o','LineStyle','none','LineWidth',2,'MarkerSize',10);
        end
    end
else
    i = plotting.iteration;
    switch plotting.model
        case 'surrogate'
            plot3(pareto_front.iter(i).surrogate(:,num_params+1),pareto_front.iter(i).surrogate(:,num_params+2),pareto_front.iter(i).surrogate(:,num_params+3),'color','r','marker','o','LineStyle','none','LineWidth',2,'MarkerSize',10);
            
        case 'truth'
            plot3(pareto_front.iter(i).truth(:,num_params+1),pareto_front.iter(i).truth(:,num_params+2),pareto_front.iter(i).truth(:,num_params+3),'color','r','marker','o','LineStyle','none','LineWidth',2,'MarkerSize',10);
    end
end
% Format axes
grid on;
set(pareto_plot,'Color','w');
set(gca,'Fontsize',plotting.afs,'Fontname',plotting.fn);
xlab = xlabel('Averaged Stress - Axial');
set(xlab,'Fontsize',plotting.afs,'Fontname',plotting.fn);
ylab = ylabel('Averaged Stress - Moment');
set(ylab,'Fontsize',plotting.afs,'Fontname',plotting.fn);
zlab = zlabel('Design Volume');
set(zlab,'Fontsize',plotting.afs,'Fontname',plotting.fn);
% Save as PDF
print(pareto_plot,'pareto3d','-dpdf');
close all
disp('done');